package io.yanglong.aiassistant.agent;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.yanglong.aiassistant.service.DocSearcher;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

@Slf4j
@Service
public class DocumentAgentImpl implements DocumentAgent {
    private final EmbeddingModel qwenEmbeddingModel;
    private final EmbeddingStore<TextSegment> redisEmbeddingStore;
    private final DocSearcher docSearcher;

    @Autowired
    public DocumentAgentImpl(EmbeddingModel qwenEmbeddingModel, EmbeddingStore<TextSegment> redisEmbeddingStore, DocSearcher qwenDocSearcher) {
        this.qwenEmbeddingModel = qwenEmbeddingModel;
        this.redisEmbeddingStore = redisEmbeddingStore;
        this.docSearcher = qwenDocSearcher;
    }

    @Override
    public String search(String userId, String userMsg) {
        return docSearcher.search(userId, userMsg);
    }

    @Override
    public boolean embeddingDocByPlain(String doc) {
        boolean success = false;
        try {
            TextSegment segment = TextSegment.from(doc);
            Embedding embedding = qwenEmbeddingModel.embed(segment).content();
            //向量化后的数据要和原文一起存入向量数据库
            redisEmbeddingStore.add(embedding, segment);
            success = true;
        } catch (Exception e) {
            log.error("Error while embedding doc,doc is {}", doc, e);
        }
        return success;
    }

    @Override
    public boolean embeddingDocByPath(String filePath) {
        return false;
    }
}
